import cv2
import numpy as np
import matplotlib.pyplot as plt
from skimage import io, img_as_float
from skimage.metrics import mean_squared_error
######################### 注册新的裁剪层方法 ##############################
# 定义一个名为CropLayer的类
class CropLayer(object):
    # 构造函数，用于初始化裁剪的起始和结束位置
    def __init__(self, params, blobs):
        self.xstart = 0
        self.xend = 0
        self.ystart = 0
        self.yend = 0

    # 该层接收两个输入。
    # 我们需要将第一个输入blob裁剪成与第二个输入blob相同的形状（保持批大小和通道数不变）
    def getMemoryShapes(self, inputs):
        # 获取输入形状和目标形状
        inputShape, targetShape = inputs[0], inputs[1]
        # 获取批大小和通道数
        batchSize, numChannels = inputShape[0], inputShape[1]
        
        height_in, width_in = inputShape[2], inputShape[3] # 获取输入的高度和宽度
        height_out, width_out = targetShape[2], targetShape[3] # 获取目标形状的高度和宽度
        
        # 居中裁剪方案
        self.ystart = (height_in - height_out) //2
        self.xstart = (width_in - width_out) //2
        
        # 左上角裁剪方案
        # self.ystart=0 
        # self.xstart=0
        
        # 计算裁剪的结束位置
        self.yend = self.ystart + height_out
        self.xend = self.xstart + width_out
        # 返回裁剪后的形状
        return [[batchSize, numChannels, height_out, width_out]]

    # 前向传播函数，进行裁剪操作
    def forward(self, inputs):
        return [inputs[0][:, :, self.ystart:self.yend, self.xstart:self.xend]]

# 使用cv2.dnn_registerLayer函数将自定义裁剪层注册到网络中
cv2.dnn_registerLayer('Crop', CropLayer)
##############################################################################


# 该函数使用预训练的 HED (Holistically-Nested Edge Detection) 模型来检测图像中的边缘。
# 通过使用 OpenCV 库加载模型并传递输入图像，该函数生成并返回二值化的边缘检测结果。
# 参数:
#   img_rgb (numpy.ndarray): 输入的 RGB 图像，图像尺寸应为 (height, width, 3)。
#   blur_window (int, 可选): 用于对输入图像和边缘检测结果进行高斯模糊的窗口大小。默认值为 5。
#   scalefactor (float, 可选): 输入图像在创建 blob 时的缩放系数。默认值为 1.0。
# 返回:
#   hed (numpy.ndarray): 二值化的边缘检测结果，图像尺寸与输入图像相同。
# 示例:
#   >>> img = cv2.imread("example.jpg")
#   >>> hed = get_hed(img)
#   >>> plt.imshow(hed)


def get_hed(img_rgb, blur_window=5, scalefactor=1.0):
    # 加载预训练的HED模型
    model = cv2.dnn.readNetFromCaffe('model\\deploy.prototxt', 'model\\hed_pretrained_bsds.caffemodel')
    
    # 对输入图像进行预处理
    IMG = 'pebbles.jpg'
    img_rgb = cv2.imread(IMG, 1)
    img_gray = cv2.imread(IMG, 0)
    blob = cv2.dnn.blobFromImage(img_rgb, scalefactor=scalefactor, size=(img_rgb.shape[1], img_rgb.shape[0]), mean=(104.00698793, 116.66876762, 122.67891434), swapRB=False, crop=False)
    
    # 通过模型进行边缘检测
    model.setInput(blob)
    hed = model.forward()
    hed = np.squeeze(hed)
    
    # 对边缘检测结果进行二值化处理
    hed = (hed > 0.5).astype(np.uint8) * 255
    
    # 对输入图像和边缘检测结果进行高斯模糊
    hed = cv2.GaussianBlur(hed, (blur_window, blur_window), 0)
    img_gray = cv2.GaussianBlur(cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY), (blur_window, blur_window), 0)
    
    return hed


# 该函数通过比较输入的灰度图像与 HED 边缘检测结果之间的均方误差 (MSE)，自动选择最佳的 Canny 边缘检测阈值参数。使用最佳阈值参数计算最终的边缘检测结果。
# 参数:
#   img_gray (numpy.ndarray): 输入的灰度图像，图像尺寸应为 (height, width)。
#   hed (numpy.ndarray): 使用 HED 模型生成的边缘检测结果，图像尺寸应与输入灰度图像相同。
# 返回:
#   edge_img (numpy.ndarray): 使用最佳阈值参数生成的 Canny 边缘检测结果。
#   best_params (tuple): 包含最佳阈值参数的元组，格式为 (best_sigma, min_mse, lower, upper)。
#       best_sigma: 为具有最低MSE的sigma值
#       min_mse: 最低的mse值
#       lower, upper: 用best_sigma进行canny时的上下限
#   plot_data (tuple): 包含用于绘制 MSE 与 sigma 值关系图的数据，格式为 (sigma_values, mse_values)。
#       sigma_values: 所有测试过的sigma_value,建议范围从0.001到1.0实验100个值
#       mse_values: 对于每一个sigma的mse值
# 示例:
# >>> img_gray = cv2.cvtColor(cv2.imread("example.jpg"), cv2.COLOR_BGR2GRAY)
# >>> hed = get_hed(cv2.imread("example.jpg"))
# >>> edge_img, best_params, plot_data = auto_canny(img_gray, hed)
# >>> plt.imshow(edge_img)

def auto_canny(img_gray, hed):

    sigma_values = np.linspace(0.001, 1.0, num=100)
    mse_values = []

    best_sigma = None
    min_mse = float('inf')
    median = np.median(img_gray)
    edge_img = None
    
    lower = 0
    upper = 0
    

    for sigma in sigma_values:
        # 使用 Canny 边缘检测算法
        median = np.median(img_gray)
        lower = int(max(0, (1.0 - sigma) * median))
        upper = int(min(255, (1.0 + sigma) * median))
        canny_img = cv2.Canny(img_gray, lower, upper)
        
        # 计算 MSE
        mse = mean_squared_error(canny_img, hed)
        mse_values.append(mse)

       

        # 更新最佳参数
        if mse < min_mse:
            min_mse = mse
            best_sigma = sigma
            edge_img = canny_img

    return edge_img, (best_sigma, min_mse, lower, upper), (sigma_values, mse_values)
